# 导入必要的库和模块
import gradio as gr
import pickle
from PIL import Image
import numpy as np

# 加载保存的KNN模型，这样我们可以使用预训练的模型进行预测
with open('best_knn_model.pkl', 'rb') as f:
    knn_model = pickle.load(f)

# 定义预测函数，这个函数将用于Gradio接口进行预测
def predict_digit(img):
    # 将图像转换为灰度图并调整尺寸为8x8（与训练数据格式一致）
    img = img.convert('L')
    img = img.resize((8, 8))
    # 转换为数组并归一化（训练数据是0-16的灰度值，这里转换为0-1范围）
    img_array = np.array(img) / 16.0
    # 展平为一维数组，符合模型输入格式
    img_flat = img_array.flatten()
    # 模型预测
    prediction = knn_model.predict([img_flat])
    return int(prediction[0])

# 创建Gradio接口，这个接口将用于用户输入和显示预测结果
interface = gr.Interface(
    fn=predict_digit,
    inputs=gr.Sketchpad(label="请绘制一个数字"),
    outputs=gr.Label(label="预测结果"),
    title="KNN手写数字识别",
    description="绘制一个数字，模型将预测它是什么数字"
)

# 启动Gradio接口，用户可以通过这个接口进行交互
interface.launch()